
#include "fun.h"
#include "hs_cuda_elemwise.cuh"

template<typename T>
struct AddFunctor {
    __device__ T operator()(T x, T y) const
    {
        return x + y;
    }
};

void elemAdd(const hs::raster::RasterData<float>& host_mat_a, const hs::raster::RasterData<float>& host_mat_b, hs::raster::RasterData<float>& host_mat_res)
{
    hscuda::raster::RasterData<float> device_mat_a;
    hscuda::raster::RasterData<float> device_mat_b;
    hscuda::raster::RasterData<float> device_mat_c;

    hscuda::raster::host2device(host_mat_a.view(), device_mat_a);
    hscuda::raster::host2device(host_mat_b.view(), device_mat_b);

    int size_x = device_mat_a.view().sizeX();
    int size_y = device_mat_a.view().sizeY();
    int N = device_mat_a.view().numel();

    // 输出分配内存
    device_mat_c.init(size_x, size_y);

    /* 使用二入一出模板函数 */
    hs::cuda::elemwise::Binary(N
        , device_mat_a.view().pData()
        , device_mat_b.view().pData()
        , AddFunctor<float>()
        , device_mat_c.view().pData()
    );
    cudaDeviceSynchronize();

    device_mat_c.toHost(host_mat_res);
}



struct FunctorMeshgridX
{

public:
    FunctorMeshgridX(int size_x, int size_y) 
        : m_size_x(size_x), m_size_y(size_y)
    { }

    __device__ int operator()(int id, int tmp) const
    {
        int yid = id / m_size_x;
        int xid = id - yid * m_size_x;
        return xid;
    }

private:
    int m_size_x;
    int m_size_y;
};


struct FunctorMeshgridY
{

public:
    FunctorMeshgridY(int size_x, int size_y) 
        : m_size_x(size_x), m_size_y(size_y)
    { }

    __device__ int operator()(int id, int tmp) const
    {
        int yid = id / m_size_x;
        return yid;
    }

private:
    int m_size_x;
    int m_size_y;
};

void meshgrid(int size_x, int size_y, hs::raster::RasterData<float>& host_mat_xx, hs::raster::RasterData<float>& host_mat_yy)
{
    int len = size_x * size_y;
    hscuda::raster::RasterData<float> device_mat_xx(size_x, size_y);
    hscuda::raster::RasterData<float> device_mat_yy(size_x, size_y);
    
    hs::cuda::elemwiseIndex::Unary(len
        , device_mat_xx.view().pData()
        , FunctorMeshgridX(size_x, size_y)
        , device_mat_xx.view().pData()
    );
    cudaDeviceSynchronize();

    hs::cuda::elemwiseIndex::Unary(len
        , device_mat_yy.view().pData()
        , FunctorMeshgridY(size_x, size_y)
        , device_mat_yy.view().pData()
    );
    cudaDeviceSynchronize();

    device_mat_xx.toHost(host_mat_xx);
    device_mat_yy.toHost(host_mat_yy);
}

